#include "types.h"
#include "defs.h"
#include "param.h"
#include "mmu.h"
#include "proc.h"
#include "spinlock.h"
#include "memlayout.h"
#define MAX_SHM_PGNUM (4) // 每个共享内存最多4页内存
#define MAX_NUMOFSHM 8
#define PRINT_DEBUG 1

struct sharemem
{
    int refcount;                  // 当前共享内存的引用数量，当引用数为0时才会真正回收
    int pagenum;                   // 占用的页数（0~4）
    void *physaddr[MAX_SHM_PGNUM]; // 对应每页的物理地址
};

struct spinlock shmlock;              // 用于互斥访问的锁
struct sharemem shmtab[MAX_NUMOFSHM]; // 整个系统最多8个共享内存

void sharememinit()
{
    initlock(&shmlock, "shmlock");         // 初始化锁
    for (int i = 0; i < MAX_NUMOFSHM; i++) // 初始化shmtab
    {
        shmtab[i].refcount = 0;
    }
#if PRINT_DEBUG
    cprintf("shm init finished.\n");
#endif
}

int shmkeyused(uint key, uint mask)
{
    if (key < 0 || 8 <= key)
    {
        return 0;
    }
    return (mask >> key) & 0x01; // 判断对应的系统共享内存是否已经启用
}

void shmaddcount(uint mask)
{
    acquire(&shmlock);
    for (int key = 0; key < MAX_NUMOFSHM; key++)
    {
        if (shmkeyused(key, mask)) // 对当前进程所有引用的共享内存计数加一
        {
            shmtab[key].refcount++;
        }
    }
    release(&shmlock);
}

// map
static int mapshm(pde_t *pgdir, uint oldshm, uint newshm, uint sz, void **physaddr)
{
    uint a;
    if (oldshm & 0xFFF || newshm & 0xFFF || oldshm > KERNBASE || newshm < sz)
    {
        return 0;
    }
    a = newshm;
    for (int i = 0; a < oldshm; a += PGSIZE, i++)
    {
        mappages(pgdir, (char *)a, PGSIZE, (uint)physaddr[i], PTE_W | PTE_U);
    }
    return newshm;
}

// 如果因为未能分配到物理页帧而失败时，调用deallocshm()将已经分配的物理页帧归还
static int deallocshm(pde_t *pgdir, uint oldshm, uint newshm)
{
    pde_t *pte;
    uint a, pa;
    if (newshm <= oldshm)
    {
        return oldshm;
    }
    a = (uint)PGROUNDDOWN(newshm - PGSIZE);
    for (; oldshm <= a; a -= PGSIZE)
    {
        pte = walkpgdir(pgdir, (char *)a, 0);
        if (pte && (*pte & PTE_P) != 0)
        {
            pa = PTE_ADDR(*pte);
            if (pa == 0)
            {
                panic("kfree");
            }
            *pte = 0;
        }
    }
    return newshm;
}

// 该方法与 allcouvm 实现基本一样
static int allocshm(pde_t *pgdir, uint oldshm, uint newshm, uint sz, void *phyaddr[MAX_SHM_PGNUM])
{
    char *mem;
    uint a;
    if (oldshm & 0xFFF || newshm & 0xFFF || oldshm > KERNBASE || newshm < sz)
    {
        return 0;
    }
    a = newshm;
    for (int i = 0; a < oldshm; a += PGSIZE, i++)
    {
        mem = kalloc(); // 分配物理页帧
        if (mem == 0)
        {
#ifdef PRINT_DEBUG
            cprintf("allocshm out of memory\n");
#endif
            deallocshm(pgdir, newshm, oldshm);
            return 0;
        }
        memset(mem, 0, PGSIZE);
        mappages(pgdir, (char *)a, PGSIZE, (uint)V2P(mem), PTE_W | PTE_U);
        phyaddr[i] = (void *)V2P(mem);
    }
    return newshm;
}

static int shmadd(uint key, uint pagenum, void *physaddr[MAX_SHM_PGNUM])
{
    if (key < 0 || key >= MAX_NUMOFSHM || pagenum < 0 || MAX_SHM_PGNUM < pagenum) // 校验参数
    {
        return -1;
    }
    shmtab[key].refcount = 1;
    shmtab[key].pagenum = pagenum;
    for (int i = 0; i < pagenum; i++)
    {
        shmtab[key].physaddr[i] = physaddr[i];
    }
    return 0;
}

// 共享内存区域映射到进程空间
void *shmgetat(uint key, uint num)
{
    pde_t *pgdir;
    void *phyaddr[MAX_SHM_PGNUM];
    uint shm = 0;
    if (key < 0 || key >= 8 || num < 0 || MAX_SHM_PGNUM < num) // 校验参数
    {
        return (void *)-1;
    }
    acquire(&shmlock);
    pgdir = proc->pgdir;
    shm = proc->shm;

    // 情况1：如果当前进程已经映射了该key的共享内存，直接返回地址
    if (proc->shmkeymask >> key & 1)
    {
        release(&shmlock);
        return proc->shmva[key];
    }

    // 情况2：如果系统还未创建此key对应的共享内存，则分配内存并映射
    if (shmtab[key].refcount == 0)
    {
        // 新增加的allocshm分配内存并映射，其原理与allocuvm相同
        shm = allocshm(pgdir, shm, shm - num * PGSIZE, proc->sz, phyaddr);
        if (shm == 0)
        {
            release(&shmlock);
            return (void *)-1;
        }
        proc->shmva[key] = (void *)shm;
        shmadd(key, num, phyaddr); // 将新内存区信息填入shmtab[MAX_NUMOFSHM]数组
    }
    else
    {
        // 情况3：如果未持有且已经系统中分配此key对应的共享内存，则直接映射
        for (int i = 0; i < num; i++)
        {
            phyaddr[i] = shmtab[key].physaddr[i];
        }
        num = shmtab[key].pagenum;
        // mapshm 方法新建映射
        if ((shm = mapshm(pgdir, shm, shm - num * PGSIZE, proc->sz, phyaddr)) == 0)
        {
            release(&shmlock);
            return (void *)-1;
        }
        proc->shmva[key] = (void *)shm;
        shmtab[key].refcount++; // 引用计数加一
    }
    proc->shm = shm;
    proc->shmkeymask |= 1 << key;
    release(&shmlock);
    return (void *)shm;
}

int shmrefcount(uint key)
{
    acquire(&shmlock);
    int count;
    count = (key < 0 || key >= MAX_NUMOFSHM) ? -1 : shmtab[key].refcount;
    release(&shmlock);
    return count;
}

// shrm() 用于共享内存的物理页帧回收，从而彻底清除共享内存区域；
// 利用 kfree() 逐个释放共享内存对应的物理页帧
static int shmrm(int key)
{
    if (key < 0 || MAX_NUMOFSHM <= key)
    {
        return -1;
    }
#ifdef PRINT_DEBUG
    cprintf("shmrm: key is %d\n", key);
#endif
    struct sharemem *shmem = &shmtab[key];
    for (int i = 0; i < shmem->pagenum; i++)
    {
        kfree((char *)P2V(shmem->physaddr[i])); // 逐个页帧回收
    }
    shmem->refcount = 0;
    return 0;
}

// 解除共享内存区域的映射
int shmrelease(pde_t *pgdir, uint shm, uint keymask)
{
#ifdef PRINT_DEBUG
    cprintf("shmrelease: shm is %x, keymask is %x.\n", shm, keymask);
#endif
    acquire(&shmlock);
    deallocshm(pgdir, shm, KERNBASE); // 释放用户空间的内存
    for (int k = 0; k < MAX_NUMOFSHM; k++)
    {
        if (shmkeyused(k, keymask))
        {
            shmtab[k].refcount--;
            if (shmtab[k].refcount == 0)
            {
                shmrm(k);
            }
        }
    }
    release(&shmlock);
    return 0;
}